#include "cartographer/rcm/particle_filter/particle_filter_2d.h"

namespace cartographer
{
    void ParticleFilterGauss::Initiate(const Eigen::Vector3d &initPose, const Eigen::Vector3d &sigma)
    {
        auto generator = std::make_shared<GaussGenerator<double, 3>>(initPose, sigma);
        Initiate([&generator] -> Particle::Ptr {
            auto particle = std::make_shared<Particle>();
            auto vector = generator->Generate();
            particle->pose = std::make_shared<Pose2d>(vector);
            return particle;
        });
    }

    void ParticleFilterGauss::Prediction(const Eigen::Vector3d &velocity, double delta_t, const Eigen::Vector3d &sigma)
    {
        auto generator = std::make_shared<GaussGenerator<double, 3>>(sigma);
        for (int i = 0; i < particle_num_; ++i)
        {
            Pose2d::Ptr pose = std::dynamic_pointer_cast<Pose2d>(particles[i]->pose);
            pose->Update(velocity, delta_t);
            generator->ResetExp(pose->data());
            auto vector = generator->Generate();
            pose->Set(vector);
        }
    }

    void ParticleFilterGauss::Write(const std::string &filename)
    {
        std::ofstream dataFile;
        dataFile.open(filename, std::ios::app);
        for (int i = 0; i < particle_num_; ++i)
        {
            auto pose = std::dynamic_pointer_cast<Pose2d>(particles[i]->pose);
            Eigen::Vector3d vector = pose->data();
            dataFile << vector(0) << " " << vector(1) << " " << vector(2) << "\n";
        }
        dataFile.close();
    }

    void ParticleFilterGauss::UpdateWeight(const Map2d::Ptr &map, const std::vector<Landmark::Ptr> &observations, Eigen::Vector3d sigma)
    {
        UpdateWeight([&](const Particle::Ptr &particle) -> double {
            double weight = 1.0;
            for (const auto &obs : observations)
            {
                auto poseInMap = obs->pose->Transform(particle->pose);
                const auto poseNear = map->SearchNearPose(poseInMap);
                double probability = poseInMap->GaussProbability(poseNear, sigma);
                weight *= probability;
            }
            return weight;
        });
    }

    void ParticleFilterGauss::Resample()
    {
        std::default_random_engine gen;
        std::discrete_distribution<int> distribution(weights.begin(), weights.end());
        std::vector<Particle::Ptr> resampled_particles;
        for (int i = 0; i < particle_num_; i++)
            resampled_particles.push_back(particles[distribution(gen)]);
        particles = resampled_particles;
    }
} // namespace cartographer